package dg;

import java.util.*;

/**
 * 决策树
 */
public class DecisionTree {
    public static void main(String[] args) {
        int items[][] = {{1, 3, 2, 1, 2},
                {1, 3, 2, 2, 2},
                {2, 3, 2, 1, 1},
                {3, 2, 2, 1, 1},
                {3, 1, 1, 1, 1},
                {3, 1, 1, 2, 2},
                {2, 1, 1, 2, 1},
                {1, 2, 2, 1, 2},
                {1, 1, 1, 1, 1},
                {3, 2, 1, 1, 1},
                {1, 2, 1, 2, 1},
                {2, 2, 2, 2, 1},
                {2, 3, 1, 1, 1},
                {3, 2, 2, 2, 2}};
        decisionTree(items);
    }

    private static void decisionTree(int[][] items) {
        /*Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(fpTree);
        while (!queue.isEmpty()) {
            TreeNode poll = queue.poll();
            queue.addAll(poll.getChildTreeNodes());
            System.out.println("当前节点 " + poll.getCurrentPoint() + "    数量  " + poll.getCount() + "    路径  "
                    + poll.getItemPathList() + "   点集  " + poll.getChildPointSet());

        }*/
        Queue<QueueNode> queueNodesQueue = new LinkedList<>();
        boolean[] visit = new boolean[items[0].length];
        DecisionTreeNode head = new DecisionTreeNode();

        QueueNode queueNode = new QueueNode(items, visit, 4);
        queueNode.setDecisionTreeNode(head);
        queueNodesQueue.offer(queueNode);
        while (!queueNodesQueue.isEmpty()) {
            QueueNode poll = queueNodesQueue.poll();
            Integer maxGain = poll.getMaxGain();
            DecisionTreeNode decisionTreeNode = poll.getDecisionTreeNode();
            Set<Integer> types = poll.getTypes();
            if (poll.isContinue()) {
                continue;
            }
            for (Integer type : types) {
                int count = poll.getArrCount(maxGain + "_" + type);
                int[][] items1 = poll.getItems();
                DecisionTreeNode node = new DecisionTreeNode();
                node.setAttrName(maxGain);
                node.setAttrValue(type);
                decisionTreeNode.addChild(node);
                visit[maxGain] = true;
                int[][] temp = new int[count][];
                Integer row = poll.getRow();
                int index = 0;
                for (int i = 0; i < row; i++) {
                    if (items1[i][maxGain] == type) {
                        temp[index] = items1[i];
                        index++;
                    }
                }
                queueNode = new QueueNode(temp, visit, 4);
                queueNode.setDecisionTreeNode(node);
                queueNodesQueue.offer(queueNode);
            }

        }
        //深度优先遍历
        Dfs(head);
    }

    private static void Dfs(DecisionTreeNode head) {
        if (head == null) {
            return;
        }
        List<DecisionTreeNode> treeNodeList = head.getTreeNodeList();
        for (DecisionTreeNode node : treeNodeList) {
            System.out.println(node.getAttrName() + "   " + node.getAttrValue());
            Dfs(node);
        }
    }
}

/**
 * 队列中保存的节点
 */
class QueueNode {
    private int[][] items;     //保存将要计算增益信息得数组
    private boolean[] visit; //保存属性是否已经纳入了决策树中
    private DecisionTreeNode decisionTreeNode;//指向将要进行计算得决策树节点
    private Set<Integer> types;//保存种类信息
    private Integer maxGain;    //信息增益最大的节点
    private Integer classifyNode; //需要分类的节点
    private Map<String, Integer> classifyMap = new HashMap<>();
    private Integer row;    //行数
    private Integer col;    //列数
    private Float mean; //期望
    List<Set<Integer>> sets = new ArrayList<>();//保存每一列得类型

    public QueueNode(int[][] items, boolean[] visit, Integer classifyNode) {
        this.items = items;
        this.visit = visit;
        this.classifyNode = classifyNode;
        //初始化
        init();
        //计算信息熵
        computeInfoGain();
        this.types = sets.get(this.maxGain);
    }


    /**
     * 计算log2
     *
     * @param N
     * @return
     */
    public float log2(double N) {
        if (N == 0) {
            return 0;
        }
        return (float) Math.log(N) / (float) Math.log(2);//Math.log的底为e
    }

    private void computeInfoGain() {
        float max = Float.NEGATIVE_INFINITY;
        for (int i = 0; i < visit.length; i++) {
            if (visit[i] || i == this.classifyNode) {
                continue;
            }
            float sum = 0F;
            Set<Integer> types = sets.get(i);
            for (Integer type : types) {
                String suffixType = i + "_" + type + "_";
                Integer count = this.classifyMap.get(i + "_" + type);
                float mean = computeMean(suffixType, count);
                sum = sum + (count / (float) this.row) * mean;
            }
            if (max < this.mean - sum) {
                max = this.mean - sum;
                this.maxGain = i;
            }
        }
    }

    private float computeMean(String prefixType, float row) {
        Set<Integer> types = sets.get(classifyNode);
        float sum = 0F;
        for (Integer type : types) {
            Integer count = classifyMap.get(prefixType + type);
            if (count == null) {
                count = 0;
            }
            sum = sum - (count / (float) row * log2(count / (float) row));
        }
        return sum;
    }

    private void init() {
        //初始化行数与列数
        this.row = this.items.length;
        this.col = this.items[0].length;
        Set<Integer> set = null;
        for (int i = 0; i < this.col; i++) {
            set = new HashSet<>();
            this.sets.add(set);
            for (int j = 0; j < this.row; j++) {
                if (this.visit[i]) {
                    continue;
                }

                set.add(this.items[j][i]);
                String key1 = i + "_" + this.items[j][i] + "_" + this.items[j][this.classifyNode];
                if (!this.classifyMap.containsKey(key1)) {
                    this.classifyMap.put(key1, 1);
                } else {
                    Integer count = this.classifyMap.get(key1);
                    count++;
                    this.classifyMap.put(key1, count);
                }

                String key2 = i + "_" + this.items[j][i];
                if (!this.classifyMap.containsKey(key2)) {
                    this.classifyMap.put(key2, 1);
                } else {
                    Integer count = this.classifyMap.get(key2);
                    count++;
                    this.classifyMap.put(key2, count);
                }
            }
        }
        this.mean = computeMean(this.classifyNode + "_", this.row);
    }

    public void setItems(int[][] items) {
        this.items = items;
    }

    public void setDecisionTreeNode(DecisionTreeNode decisionTreeNode) {
        this.decisionTreeNode = decisionTreeNode;
    }

    public DecisionTreeNode getDecisionTreeNode() {
        return decisionTreeNode;
    }

    public Integer getMaxGain() {
        return maxGain;
    }

    public Set<Integer> getTypes() {
        return types;
    }

    public int getArrCount(String key) {
        return this.classifyMap.get(key);
    }

    public Integer getRow() {
        return row;
    }

    public int[][] getItems() {
        return items;
    }

    public boolean isContinue() {
        return this.sets.get(classifyNode).size() == 1;
    }
}

/**
 * 决策树节点
 */
class DecisionTreeNode {
    private Integer attrName;    //属性名称
    private Integer attrValue;   //属性值
    private List<DecisionTreeNode> treeNodeList = new ArrayList<>();

    public void setAttrName(Integer attrName) {
        this.attrName = attrName;
    }

    public void setAttrValue(Integer attrValue) {
        this.attrValue = attrValue;
    }

    public void addChild(DecisionTreeNode decisionTreeNode) {
        treeNodeList.add(decisionTreeNode);
    }

    public boolean hasChild() {
        return treeNodeList.isEmpty();
    }

    public Integer getAttrName() {
        return attrName;
    }

    public Integer getAttrValue() {
        return attrValue;
    }

    public List<DecisionTreeNode> getTreeNodeList() {
        return treeNodeList;
    }
}
